#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
微调模型推理脚本
用于测试微调后的模型效果
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import argparse
import os

def load_model(base_model_name, lora_path):
    """
    加载基础模型和LoRA权重
    
    Args:
        base_model_name (str): 基础模型名称
        lora_path (str): LoRA权重路径
    
    Returns:
        tuple: (model, tokenizer)
    """
    print(f"正在加载基础模型: {base_model_name}")
    
    # 检测设备类型 - 为了避免MPS内存问题，优先使用CPU
    if torch.cuda.is_available():
        device = "cuda"
        torch_dtype = torch.bfloat16
        print("使用 CUDA 设备")
    else:
        # 暂时使用CPU避免MPS内存限制问题
        device = "cpu"
        torch_dtype = torch.float32
        print("使用 CPU 设备 (为避免MPS内存问题)")
        if torch.backends.mps.is_available():
            print("注意: 检测到MPS设备，但为避免内存问题使用CPU")
    
    # 加载分词器
    tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_fast=True)
    
    # 加载基础模型 - 根据设备类型调整参数
    if device == "cuda":
        # CUDA 设备可以使用量化
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            torch_dtype=torch_dtype,
            device_map="auto",
            load_in_4bit=True
        )
    else:
        # CPU 和 MPS 设备不使用量化，启用低内存模式
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            torch_dtype=torch_dtype,
            device_map=None,  # 禁用自动设备映射避免内存问题
            low_cpu_mem_usage=True,
            trust_remote_code=True
        )
        # 手动移动模型到设备
        base_model = base_model.to(device)
    
    # 加载LoRA权重
    if os.path.exists(lora_path):
        print(f"正在加载LoRA权重: {lora_path}")
        model = PeftModel.from_pretrained(base_model, lora_path)
    else:
        print(f"警告: LoRA权重路径不存在 {lora_path}，使用原始模型")
        model = base_model
    
    return model, tokenizer

def generate_response(model, tokenizer, instruction, input_text="", max_length=512, temperature=0.7):
    """
    生成模型回复
    
    Args:
        model: 加载的模型
        tokenizer: 分词器
        instruction (str): 指令
        input_text (str): 输入文本（可选）
        max_length (int): 最大生成长度
        temperature (float): 生成温度
    
    Returns:
        str: 生成的回复
    """
    # 构建提示词
    if input_text:
        prompt = f"指令: {instruction}\n输入: {input_text}\n回答:"
    else:
        prompt = f"指令: {instruction}\n回答:"
    
    # 编码输入并移动到正确的设备
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
    # 生成回复
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            temperature=temperature,
            do_sample=True,
            top_p=0.9,
            top_k=50,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id
        )
    
    # 解码输出
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 提取回答部分
    if "回答:" in response:
        answer = response.split("回答:")[-1].strip()
    else:
        answer = response.strip()
    
    return answer

def interactive_chat(model, tokenizer):
    """
    交互式对话模式
    
    Args:
        model: 加载的模型
        tokenizer: 分词器
    """
    print("\n=== 交互式对话模式 ===")
    print("输入 'quit' 或 'exit' 退出")
    print("输入格式: 指令 [| 输入内容]")
    print("示例: 翻译以下句子 | Hello world")
    print("="*50)
    
    while True:
        try:
            user_input = input("\n用户: ").strip()
            
            if user_input.lower() in ['quit', 'exit', '退出']:
                print("再见！")
                break
            
            if not user_input:
                continue
            
            # 解析输入
            if '|' in user_input:
                instruction, input_text = user_input.split('|', 1)
                instruction = instruction.strip()
                input_text = input_text.strip()
            else:
                instruction = user_input
                input_text = ""
            
            # 生成回复
            print("\n模型正在思考...")
            response = generate_response(model, tokenizer, instruction, input_text)
            print(f"\n助手: {response}")
            
        except KeyboardInterrupt:
            print("\n\n再见！")
            break
        except Exception as e:
            print(f"\n错误: {e}")

def main():
    parser = argparse.ArgumentParser(description="微调模型推理脚本")
    parser.add_argument(
        "--base_model", 
        type=str, 
        default="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
        help="基础模型名称"
    )
    parser.add_argument(
        "--lora_path", 
        type=str, 
        default="./models/qwen15b-lora",
        help="LoRA权重路径"
    )
    parser.add_argument(
        "--instruction", 
        type=str, 
        help="单次推理的指令"
    )
    parser.add_argument(
        "--input", 
        type=str, 
        default="",
        help="单次推理的输入"
    )
    parser.add_argument(
        "--interactive", 
        action="store_true",
        help="启动交互式对话模式"
    )
    
    args = parser.parse_args()
    
    # 加载模型
    model, tokenizer = load_model(args.base_model, args.lora_path)
    
    if args.interactive:
        # 交互式模式
        interactive_chat(model, tokenizer)
    elif args.instruction:
        # 单次推理
        response = generate_response(model, tokenizer, args.instruction, args.input)
        print(f"\n指令: {args.instruction}")
        if args.input:
            print(f"输入: {args.input}")
        print(f"回答: {response}")
    else:
        # 默认测试
        test_cases = [
            {"instruction": "解释什么是人工智能", "input": ""},
            {"instruction": "翻译以下句子", "input": "Hello, how are you?"},
            {"instruction": "写一首关于秋天的短诗", "input": ""}
        ]
        
        print("\n=== 模型测试 ===")
        for i, case in enumerate(test_cases, 1):
            print(f"\n测试 {i}:")
            print(f"指令: {case['instruction']}")
            if case['input']:
                print(f"输入: {case['input']}")
            
            response = generate_response(model, tokenizer, case['instruction'], case['input'])
            print(f"回答: {response}")
            print("-" * 50)

if __name__ == "__main__":
    main()